import json
import re

from Agent.Agent import Agent
from llm import llm, super_eval

from Agent.NavigateAgent import navigate_agent
from Agent.CodeAgent import CodeAgentConfig, CodeAgent
from Agent.DeepSeaAgent2 import deep_sea_agent
from Agent.RagAgent import rag_agent
from get_sample import get_flow_sample

prompt = '''

### 工作机器人职责说明

1. **任务状态机器人**：负责判断`深海作业A`的具体动作时间和次数等信息。
   `深海作业A` 为下放探测器任务。
   - **下放阶段**：包括以下动作：A架开机, ON DP, 折臂吊车开机, 小艇检查完毕, 小艇入水, 征服者起吊, 征服者入水, 缆绳解除, A架摆回, 小艇落座, 折臂吊车关机, A架关机, OFF DP
   - **回收阶段**：包括以下动作：A架开机, 折臂吊车开机, A架摆出, 小艇检查完毕, 小艇入水, 缆绳挂妥, 征服者出水, 小艇落座, 折臂吊车关机, 征服者落座, A架关机

   示例问题：请指出2024/09/01征服者起吊以及征服者入水发生的时间（以XX:XX输出，如有多个动作请使用英文逗号隔开；如果没有进行深海作业，回答N）。

2. **航行状态机器人**：负责判断停泊状态、航渡状态、动力定位状态和伴航状态的时间和次数等信息。

   示例问题：请统计2024年6月4日处于航渡状态的时长（单位：分钟，如果没有则计为0）。

3. **数据处理机器人**：负责进行数据查找、功率计算、判断哪些设备在运行等任务,拥有编码解决问题的能力，适合所有计算，可以解决稍微复杂的任务。

   示例问题：计算2020-05-21七点到九点甲板机械的总做功（单位：kWh，保留两位小数）。
           请指出2024年5月23日23:00至24:00进行的作业使用的甲板机器设备有哪些？
           假设柴油的密度为0.8448kg/L，柴油热值为42.6MJ/kg，请计算2024/8/23 0:00 ~ 2024/8/25 0:00的理论发电量（单位化成kWh，保留2位小数）？

4. **设备参数机器人**：负责确认各个部件的运作参数范围是否报警等。

   示例问题：停泊发电机组额定转速多少？
           停泊发电机组转速为1800RPM会发生什么？
           假设某一时刻应急发电机组的滑油温度为200℃，是否会触发报警？（Y：是；N：否。请回答Y或N）

### 思考方式

你会先进行思考规划，然后一步步执行，步骤最多不超过3个步骤。

### 使用格式

你可以使用以下格式调度不同的机器人：
```json
{"role":"机器人名称","question":"问题"}
```

### 示例
user:2024/8/23 10:17什么设备在进行什么动作？？
assistant:问题只需要调用任务状态机器人即可解决：
```json
{"role":"任务状态机器人","2024/8/23 10:17什么设备在进行什么动作？"}
```

user:假设柴油的密度为0.8448kg/L，柴油热值为42.6MJ/kg，请计算2024/8/23 0:00 ~ 2024/8/25 0:00的理论发电量（单位化成kWh，保留2位小数）？
assistant:问题只需要调用数据处理机器人即可解决：
```json
{"role":"数据处理机器人","question":"假设柴油的密度为0.8448kg/L，柴油热值为42.6MJ/kg，请计算2024/8/23 0:00 ~ 2024/8/25 0:00的理论发电量（单位化成kWh，保留2位小数）？"}
```
user:请统计2024-06-14 00:00到2024-06-14 24:00一号主推进器的功率最大值
assistant:问题只需要调用数据处理机器人即可解决：
```json
{"role":"数据处理机器人","question":"请统计2024-06-14 00:00到2024-06-14 24:00一号主推进器的功率最大值"}
```
user:在2024年8月24日深海作业A中，请输出布放和回收的开始和结束时间（以折臂吊车开机为开始，折臂吊车关机为结束，时间以24小时制HH:MM格式按发生顺序输出）
assistant:问题只需要调用数据处理机器人即可解决：
```json
{"role":"数据处理机器人","question":"请统计2024-06-14 00:00到2024-06-14 24:00一号主推进器的功率最大值"}
```

### 示例2

如果需要多次调用，请使用多轮对话。示例如下：

user:在2024年6月6日，请输出征服者起吊和征服者入水的时间（以XX:XX格式输出，时间为24小时制，时间补零，如05:03），以及在这段时间内推进器的总做功（单位为kWh，结果保留两位小数）。

assistant:我首先需要调用任务状态机器人获取时间信息，然后调用数据处理机器人进行功率计算，首先进行第一步：
```json
{"role":"任务状态机器人","question":"在2024年6月6日，请输出征服者起吊和征服者入水的时间（以XX:XX格式输出，时间为24小时制，时间补零，如05:03）"}
```
user:任务状态机器人:XX:XX-XX:XX
assistant:接下来我需要根据返回的时间，调用数据处理机器人进行功率计算：
```json
{"role":"数据处理机器人","question":"在2024年6月6日,XX:XX-XX:XX推进器的总做功多少？（单位为kWh，结果保留两位小数）"}
```

### 示例3

user:请统计2024年6月14日处于伴航状态的总时长（单位：分钟，如果没有则计为0），以及伴航状态中一号、二号、三号和四号柴油发电机各自的运行时长（单位：分钟）。
assistant:我首先需要调用航行状态机器人获取时间信息，得到了时间信息后，我需要通过数据处理机器人计算时间差，并且获取一号、二号、三号和四号柴油发电机各自的运行时长。：
```json
{"role":"航行状态机器人","question":"请统计2024年6月14日处于伴航状态的开始结束时间"}
```
user:XX:XX-XX:XX
assistant:得到了时间信息，我需要通过数据处理机器人计算时间差。
计算时间差:
```json
{"role":"数据处理机器人","question":"请计算XX:XX-XX:XX经过了多少分钟"}
```
user:xxx分钟
assistant:最后需要获取一号、二号、三号和四号柴油发电机各自的运行时长:
```json
{"role":"数据处理机器人","question":"请统计XX:XX-XX:XX一号、二号、三号和四号柴油发电机各自的运行时长"}
```
示例4：展示可以合并的部分和要求必须完整
user:2024年5月23日，从征服者落座开始至当日作业结束，请计算A架右舷摆动的次数。（以整数输出，A架摆动的定义如下：A架右舷同一方向上摆动超过10°即可算作一次摆动）
assistant:我首先需要确定2024年5月23日，从征服者落座开始至当日作业结束的时间，需要调用任务状态机器人，然后需要调用数据处理机器人计算A架右舷摆动的次数，调用数据处理机器人的时候要将问题描述完整，首先是第一步：
```json
{"role":"任务状态机器人","question":"2024年5月23日，从征服者落座和当日作业结束的时间是多少？"}
```
user:XX:XX-XX:XX
assistant:我已经获取了时间，接下来我需要诉数据处理机器人计算A架右舷摆动的次数，并将定义告诉数据处理机器人：
```json
{"role":"数据处理机器人","question":"2024年5月23日XX:XX-XX:XX，A架右舷摆动的次数？（以整数输出，A架摆动的定义如下：A架右舷同一方向上摆动超过10°即可算作一次摆动）"}
```
警告：你禁止自己解题，只能调用机器人得到答案后进行提取复述。
备注:当前为测试阶段，所有题目都可以通过调用机器人的方式解答。
你一次只会输出一个json指令，并且你不会自己进行计算。
'''

import threading
import json
import re


class FlowAgent(Agent):

    def __init__(self, name,
                 description,
                 system_prompt,
                 agent_type,
                 llm_engine,
                 sample=None,
                 agents=None,
                 rag=None):
        super().__init__(name,
                         description,
                         system_prompt,
                         agent_type,
                         llm_engine,
                         agents,
                         rag)
        self.code_agent = CodeAgent(**CodeAgentConfig)
        self.sample = sample

    def pre_process(self, question):
        question = question.strip().replace('：', ':').replace("'", '"')
        return question

    def run_with_timeout(self, question, timeout=999):
        self.result = 'N5192038476'

        def target():
            self.result = self._run(question)

        thread = threading.Thread(target=target)
        thread.start()
        thread.join(timeout)
        if thread.is_alive():
            try:
                self.code_agent.kernel.shutdown()
            except:
                ...
            raise TimeoutError("The function run exceeded the time limit of {} seconds".format(timeout))
        return self.result

    def _run(self, question):
        question = self.pre_process(question)
        self.messages = [{"role": "system", "content": self.system_prompt}]
        sample_msg = self.sample.get_sample(question, 1)
        self.messages.extend(sample_msg[0])
        self.messages.append({"role": "user", "content": '以下是新问题:' + question + "\n请输出第一个指令"})

        for _ in range(7):
            answer = self.llm_engine(self.messages)
            if '```json' in answer:
                json_list = re.findall('```json(.*?)```', answer, re.DOTALL)

                for s in json_list:
                    if 'role' not in s:
                        return answer

                    answer_json = json.loads(s.strip())
                    if answer_json['role'] == '任务状态机器人':
                        task_answer = deep_sea_agent.run(answer_json['question'])
                    elif answer_json['role'] == '航行状态机器人':
                        task_answer = navigate_agent.run(answer_json['question'])
                    elif answer_json['role'] == '数据处理机器人':
                        task_answer = self.code_agent.run(answer_json['question'])
                    elif answer_json['role'] == '设备参数机器人':
                        task_answer = rag_agent.run(answer_json['question'])
                    else:
                        task_answer = "机器人调用错误，只有：任务状态机器人，航行状态机器人，数据处理机器人，设备参数机器人 可以选择"

                    role = answer_json['role']
                    task_answer = task_answer.split('</think>')[-1]
                    self.messages.append(
                        {'role': 'user', 'content': f"{role}:`{task_answer}`\n以上是运行结果，请基于结果确定是继"
                                                    f"续调用还是输出答案。如果输出答案，请符合题目要求并尽可能简短，"
                                                    f"只输出纯答案+单位，数字不要有千分位符号。"
                                                    f"如果调用机器人，请仅输出一个指令。"})
            else:
                with open('./sample/flow.jsonl', 'a', encoding='utf-8') as fw:
                    fw.write(
                        json.dumps({"question": question, "messages": self.messages[1:]}, ensure_ascii=False)
                    )
                    fw.write('\n')

                return answer

        self.messages.append({'role': 'user', 'content': f"基于历史对话，请尽可能回答问题:{question}"})
        return self.llm_engine(self.messages)

    def run(self, question):
        try:
            res = self.run_with_timeout(question)
            try:
                self.code_agent.kernel.shutdown()
            except:
                ...
            return res
        except TimeoutError as e:
            return str(e)


FlowAgentConfig = {"name": 'xxx',
                   "description": 'xxx',
                   'system_prompt': prompt,
                   'agent_type': 'plan-action',
                   'llm_engine': llm,
                   'sample': get_flow_sample(),
                   'rag': None,
                   'agents': None}
